import {Annotation, Command, START, StateGraph} from "@langchain/langgraph";
import {saveImages} from "../tool";

const StateAnnotation = Annotation.Root({
    foo: Annotation<string>
})

const nodeA = async (_state: typeof StateAnnotation.State) => {
    console.log("Called A")

    const goto = Math.random() > .5? 'nodeB': "nodeC";

    return new Command({
        update: {
            foo: 'a'
        },
        goto
    })
}

const nodeB = async (state: typeof StateAnnotation.State) => {
    console.log("Called B")
    return {
        foo: state.foo + '|b'
    }
}
const nodeC = async (state: typeof StateAnnotation.State) => {
    console.log("Called C")
    return {
        foo: state.foo + '|c'
    }
}

const graph = new StateGraph(StateAnnotation)
.addNode("nodeA", nodeA, {
    ends: ['nodeB', 'nodeC']
}).addNode('nodeB', nodeB)
.addNode('nodeC', nodeC)
.addEdge(START, 'nodeA')
.compile()

await saveImages(graph, 'command.png')

console.log(await graph.invoke({foo: ""}))